#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# @Project : chat_model 
# @File    : minio_util.py
# @IDE     : PyCharm 
# @Author  :ZH
# @Time    : 2025/1/13 9:15
__all__ = ['MinioUtil']


import io
import json

from minio import Minio


class MinioUtil:

    def __init__(self, endpoint, access_key, secret_key):
        self.endpoint = endpoint
        self.access_key = access_key
        self.secret_key = secret_key
        self.client = None
        self.policy = {
            "Version": "2012-10-17",
            "Statement": [
                {
                    "Effect": "Allow",
                    "Principal": {"AWS": "*"},
                    "Action": [
                        "s3:GetBucketLocation",
                        "s3:ListBucket",
                        "s3:ListBucketMultipartUploads",
                    ],
                    "Resource": "arn:aws:s3:::my-bucket",
                },
                {
                    "Effect": "Allow",
                    "Principal": {"AWS": "*"},
                    "Action": [
                        "s3:GetObject",
                        "s3:PutObject",
                        "s3:DeleteObject",
                        "s3:ListMultipartUploadParts",
                        "s3:AbortMultipartUpload",
                    ],
                    "Resource": "arn:aws:s3:::my-bucket/*",
                },
            ],
        }

    def create_client(self):
        if not self.client:
            self.client = Minio(
                self.endpoint,
                access_key=self.access_key,
                secret_key=self.secret_key,
                secure=False
            )

    def create_bucket(self, bucket_name: str):
        self.create_client()
        if not self.client.bucket_exists(bucket_name):
            self.client.make_bucket(bucket_name)
            self.client.delete_bucket_policy(bucket_name)
            self.client.set_bucket_policy(bucket_name, json.dumps(self.policy).replace('my-bucket', bucket_name))

    def remove_bucket(self, bucket_name: str):
        self.create_client()
        if self.client.bucket_exists(bucket_name):
            self.client.remove_bucket(bucket_name)

    def get_object(self, bucket_name: str, object_name: str):
        self.create_client()
        if self.client.bucket_exists(bucket_name):
            try:
                # 获取文件对象
                response = self.client.get_object(bucket_name, object_name)
                # 读取文件内容
                file_data = response.read()
                return file_data
            except Exception as e:
                print(f"获取文件流时出错: {e}")
                return None
            finally:
                response.close()
                response.release_conn()
        else:
            return None

    def put_object(self, bucket_name: str, object_name: str, data):
        self.create_client()
        if self.client.bucket_exists(bucket_name):
            data_stream = io.BytesIO(data)
            self.client.put_object(bucket_name, object_name, data_stream, length=len(data))
            return object_name
        else:
            return None

    def remove_object(self, bucket_name: str, object_name: str):
        self.create_client()
        if self.client.bucket_exists(bucket_name):
            self.client.remove_object(bucket_name, object_name)
        else:
            return None


# if __name__ == '__main__':
#     minio_util = MinioUtil(
#         endpoint='6.tcp.cpolar.cn:13710', access_key='zhangheng',
#         secret_key='zhangheng0906'
#     )
#     minio_util.create_client()
#     # minio_util.close_client()
#     # minio_util.remove_bucket(bucket_name='chat-model')
#     minio_util.create_bucket(bucket_name='chat-model')
#     with open(r'D:\data\a.png', 'rb') as f:
#         minio_util.put_object(bucket_name='chat-model', object_name='test/a.png', data=f.read())
#
#     with open('./a.png', 'wb') as f:
#         # response = minio_util.get_object(bucket_name='chat-model', object_name='test/a.png')
#         # for chunk in response.stream(32 * 1024):  # 分块读取，避免内存占用过大
#         #     f.write(chunk)
#         f.write(minio_util.get_object(bucket_name='chat-model', object_name='test/a.png'))
